"""
Create at 20250331

定义生成SQL的核心逻辑。
输入 Query
输出 SQL
"""

from nl2sql.datasource.base import BaseDatasource
from typing import List, Optional
from nl2sql.prompt.generate_sql import sql_prompt
from nl2sql.model.llm import LLM
from nl2sql.tools.extract_block import extract_sql


class SQLGenerator:
    def __init__(self,
                 datasource: BaseDatasource,
                 model=None,
                 api_key=None,
                 base_url=None) -> None:
        """
        初始化函数
        :param datasource: 数据源的实例
        :param model: 模型名称
        :param api_key: 模型apikey
        :param base_url: 模型的调用地址
        """
        self.datasource: Optional[BaseDatasource] = datasource
        self.prompt: str = sql_prompt
        self.model = model or "qwen-turbo"
        self.api_key = api_key or "sk-34f5c792513c423f90b404d28b070f1f"
        self.base_url = base_url or "https://dashscope.aliyuncs.com/compatible-mode/v1"
        self.llm = LLM(model_name=self.model,
                       api_key=self.api_key,
                       base_url=self.base_url)
        pass

    def generate(self, context: str) -> str:
        """
        生成SQL
        :param context: 上下文信息，包含了用户问题，以及其他的信息
        :return: str 对应的SQL
        """
        data_info: dict = self.datasource.get_data_info()
        prompt: str = self.prompt.format(data_info=data_info,
                                         query=context)
        print('=='*20)
        print(f'调试专用: {prompt}')
        print('=='*20)
        # 调用大模型进行推理
        resp = self.llm.chat(prompt=prompt)
        sql: str = extract_sql(text=resp)
        return sql


if __name__ == '__main__':
    from nl2sql.datasource.excel_datasource import ExcelDatasource

    excel_path = "D:/PythonProject_airline/nl2sql_aline/examples/航空公司航班信息表.xlsx"
    excel_ds = ExcelDatasource(source_path=excel_path, auto_create_sqlite=True)
    data_info = excel_ds.get_data_info()

    excel_ds = ExcelDatasource(source_path=excel_path, auto_create_sqlite=True)
    sg = SQLGenerator(datasource=excel_ds,)
    sql = sg.generate(context='查一下平均油耗大于20升的航空公司里面且是中国的航司，他们的平均载客量与最大载货量之比最小的是哪家航司？并告诉我它们所有的航线')
    print(sql)

    # 4. 提取 SQLs
    sql_query = sql
    if not sql_query:
        print("Error: No valid SQL found in the response.")


    # 替换表名为 'data'，与 ExcelDatasource 默认表名一致
    sql_query = sql_query.replace("df", "data")
    print(f"\nExtracted SQL:\n{sql_query}")

    # 5. 执行 SQL 并输出结果
    try:
        result = excel_ds.execute_sql(sql_query)
        print("\nQuery Result:")
        print(result)
    except Exception as e:
        print(f"Error executing SQL: {str(e)}")

    # 6. 关闭数据库连接
    excel_ds.close_connection()


